# -*- coding: utf-8 -*-
# 标准库导入
import sys
import os

# 第三方库导入
from PyQt5.QtCore import Qt, pyqtSignal
from PyQt5.QtWidgets import (
    QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
    QLabel, QLineEdit, QPushButton, QTableWidget, QTableWidgetItem,
    QHeaderView, QMessageBox, QStackedWidget, QComboBox, QFormLayout,
    QDialog, QDialogButtonBox, QStyleFactory, QGraphicsDropShadowEffect,
    QSizePolicy  # 添加尺寸策略类
)
from PyQt5.QtGui import QFont, QIcon, QPixmap, QColor, QDoubleValidator, QIntValidator
from tqsdk import TqApi, TqAuth, TqAccount, TqSim

# 本地模块导入
from db_manager import (
    connect_db,
    get_all_products,
    insert_product,
    update_product,
    delete_product,
    get_monitored_products,
    delete_product_by_name,
    get_product_by_id
)

# 在类内新增以下方法
class CapitalManager(QMainWindow):
    update_signal = pyqtSignal(str)
    connection_signal = pyqtSignal(str, bool)

    def _connect_thread(self):
        try:
            # 修正经纪商代码为官方格式（关键修复）
            self.tq_account = TqAccount("H徽商期货", "155831", "alibaba999")  
            self.api = TqApi(account=self.tq_account,
                           auth=TqAuth("szlwy68", "alibaba999"),
                           _stock=False)
            self.macc=self.api.get_account()
            print("已经成功连接天勤账户")
            print("当前账户可用资金：",self.macc.available)
            
            # 修正f-string语法错误
            self.update_signal.emit(f"{self.macc.available:.2f}")  # 新增
            
            # 发射连接成功信号
            self.connection_signal.emit("已连接 (点击断开)", True)
            self.update_capital()  # 启动资金更新线程
            
        except Exception as e:
            self.connection_signal.emit(f"连接失败: {str(e)}", False)

    def handle_connection(self, text, success):
        """处理连接状态更新"""
        self.connect_btn.setText(text)
        self.connect_btn.setEnabled(True)
        if success:
            # 断开原有连接并绑定断开方法
            try:
                self.connect_btn.clicked.disconnect()
            except TypeError:
                pass
            self.connect_btn.clicked.connect(lambda: self.disconnect_tq())
            QApplication.processEvents()
            # 删除下面这行已移除的方法调用
            # self.update_capital()  # <<< 这行需要删除
        else:
            # 重新绑定连接方法（关键修复）
            try:
                self.connect_btn.clicked.disconnect()
            except TypeError:
                pass
            self.connect_btn.clicked.connect(self.connect_tq)

    def connect_tq(self):
        """天勤账户连接主方法"""
        try:
            self.connect_btn.setEnabled(False)
            self.connect_btn.setText("连接中...")
            
            from threading import Thread
            Thread(target=self._connect_thread, daemon=True).start()
            
        except Exception as e:
            self.connect_btn.setEnabled(True)
            QMessageBox.critical(self, "错误", f"启动连接失败: {str(e)}")

    def disconnect_tq(self):  # 新增断开连接方法
        """断开天勤账户连接"""
        try:
            if self.api:
                self.api.close()
                self.api = None
                self.tq_account = None
                self.connect_btn.setText("连接天勤账户")
                self.capital_input.clear()
                QMessageBox.information(self, "断开成功", "已断开天勤账户连接")
                
                # 重新绑定连接方法
                self.connect_btn.clicked.disconnect()
                self.connect_btn.clicked.connect(self.connect_tq)
                
        except Exception as e:
            QMessageBox.critical(self, "断开错误", f"断开连接时发生错误: {str(e)}")

    def closeEvent(self, event): 
        self.disconnect_tq()
        event.accept()

    def _connect_thread(self):
        try:
            self.tq_account = TqAccount("H徽商期货", "155831", "alibaba999")  
            self.api = TqApi(account=self.tq_account,
                           auth=TqAuth("szlwy68", "alibaba999"),
                           _stock=False)
            
            self.macc=self.api.get_account()
            print("已经成功连接天勤账户")
            print("当前账户可用资金：",self.macc.available)
            
            self.update_signal.emit(f"{self.macc.available:.2f}") 
            self.connection_signal.emit("已连接 (点击断开)", True)
            
        except Exception as e:
            self.connection_signal.emit(f"连接失败: {str(e)}", False)

    def init_auto_sl_ui(self):
        """自动止盈止损页面初始化（补充缺失的return）"""
        widget = QWidget()
        layout = QVBoxLayout(widget)
        layout.setContentsMargins(15, 15, 15, 15)

        # 策略控制按钮组
        btn_group = QHBoxLayout()
        self.strategy_switch_btn = QPushButton("启动策略止盈止损交易")
        self.open_strategy_btn = QPushButton("启动策略开仓")
        self.open_qty_input = QLineEdit("1")
        self.open_qty_input.setValidator(QIntValidator(1, 1000))
        
        # 按钮样式设置
        for btn in [self.strategy_switch_btn, self.open_strategy_btn]:
            btn.setCheckable(True)
            btn.setStyleSheet("QPushButton {padding: 8px; min-width: 120px;}")

        btn_group.addWidget(self.strategy_switch_btn)
        btn_group.addWidget(QLabel("开仓量:"))
        btn_group.addWidget(self.open_qty_input)
        btn_group.addWidget(self.open_strategy_btn)
        layout.addLayout(btn_group)

        # 监控品种表格
        self.monitor_table = QTableWidget()
        self.monitor_table.setColumnCount(6)
        self.monitor_table.setHorizontalHeaderLabels(["品种", "方向", "持仓量", "当前价", "浮动盈亏", "策略状态"])
        self.monitor_table.horizontalHeader().setSectionResizeMode(QHeaderView.Stretch)
        layout.addWidget(self.monitor_table)

        # 信号连接
        self.strategy_switch_btn.clicked.connect(self.toggle_strategy)
        self.open_strategy_btn.clicked.connect(self.toggle_open_strategy)
        
        return widget  # 确保返回创建好的widget

    def toggle_strategy(self):
        """切换止盈止损策略状态"""
        if self.strategy_switch_btn.isChecked():
            self.start_strategy()
        else:
            self.stop_strategy()

    def start_strategy(self):
        """启动策略核心逻辑"""
        try:
            if not self.api:
                raise RuntimeError("请先连接天勤账户")
            
            # 初始化策略参数
            self.strategy_running = True
            self.strategy_switch_btn.setText("停止策略止盈止损交易")
            self.open_strategy_btn.setEnabled(True)
            
            # 创建策略线程
            from threading import Thread
            self.strategy_thread = Thread(target=self.run_strategy, daemon=True)
            self.strategy_thread.start()
            
        except Exception as e:
            QMessageBox.critical(self, "错误", f"启动策略失败: {str(e)}")

    def run_strategy(self):
        """策略执行核心逻辑（参考SMA策略文件）"""
        while self.strategy_running:
            try:
                # 获取监控品种
                conn = connect_db()
                products = get_monitored_products(conn)
                conn.close()
                
                # 获取实时行情数据
                quotes = self.api.get_quote_list()
                
                for product in products:
                    symbol = product[2]  # 交易所代码
                    quote = next((q for q in quotes if q.symbol == symbol), None)
                    
                    if quote:
                        # 获取持仓数据
                        position = self.api.get_position(symbol)
                        # 获取K线数据（参考SMA策略）
                        klines = self.api.get_kline_serial(symbol, 900, 50)  # 15分钟线
                        
                        # 计算CPX均线（示例逻辑，需根据实际策略调整）
                        cpx_ma = sum(klines.close[-25:]) / 25 if len(klines) >=25 else 0
                        
                        # 策略逻辑（关键部分）
                        if position.pos_long > 0:  # 持有多单
                            if klines.close[-2] > cpx_ma and klines.close[-1] < cpx_ma:
                                self.api.insert_order(symbol, "SELL", "CLOSE", position.pos_long)
                        elif position.pos_short > 0:  # 持有空单
                            if klines.close[-2] < cpx_ma and klines.close[-1] > cpx_ma:
                                self.api.insert_order(symbol, "BUY", "CLOSE", position.pos_short)
                                
                self.api.wait_update()
            except Exception as e:
                print(f"策略执行异常: {str(e)}")
                break

    def toggle_open_strategy(self):
        """切换开仓策略状态"""
        if self.open_strategy_btn.isChecked():
            self.start_open_strategy()
        else:
            self.stop_open_strategy()

    def start_open_strategy(self):
        """启动开仓策略"""
        self.open_strategy_running = True
        self.open_strategy_btn.setText("停止策略开仓")
        self.open_qty_input.setEnabled(False)
        
        # 创建开仓策略线程
        from threading import Thread
        self.open_strategy_thread = Thread(target=self.run_open_strategy, daemon=True)
        self.open_strategy_thread.start()

    def run_open_strategy(self):
        """开仓策略逻辑（参考SMA_Open01.py）"""
        while self.open_strategy_running:
            try:
                # 获取监控品种
                conn = connect_db()
                products = get_monitored_products(conn)
                conn.close()
                
                # 获取实时行情
                quotes = self.api.get_quote_list()
                klines = {p[2]: self.api.get_kline_serial(p[2], 900, 50) for p in products}
                
                for product in products:
                    symbol = product[2]
                    # 计算策略信号（示例逻辑）
                    cpx_ma = sum(klines[symbol].close[-25:])/25 if len(klines[symbol])>=25 else 0
                    
                    # 开仓条件判断
                    if klines[symbol].close[-1] > cpx_ma:
                        # 平空开多
                        position = self.api.get_position(symbol)
                        if position.pos_short > 0:
                            self.api.insert_order(symbol, "BUY", "CLOSE", position.pos_short)
                        self.api.insert_order(symbol, "BUY", "OPEN", int(self.open_qty_input.text()))
                    elif klines[symbol].close[-1] < cpx_ma:
                        # 平多开空
                        position = self.api.get_position(symbol)
                        if position.pos_long > 0:
                            self.api.insert_order(symbol, "SELL", "CLOSE", position.pos_long) 
                        self.api.insert_order(symbol, "SELL", "OPEN", int(self.open_qty_input.text()))
                        
                self.api.wait_update()
            except Exception as e:
                print(f"开仓策略异常: {str(e)}")
                break

    def init_data(self):
        """初始化数据（补全实现）"""
        self.products = []
        self.current_product = None
        self.pages = []  # 新增页面列表
        self.current_page_index = 0

    def init_ui(self):
        """完整的主界面初始化"""
        main_widget = QWidget()
        main_layout = QVBoxLayout(main_widget)
        
        # +++ 添加账户连接控件 +++
        connect_layout = QHBoxLayout()
        self.connect_btn = QPushButton("连接天勤账户")
        self.capital_input = QLineEdit(placeholderText="可用资金")
        self.capital_input.setReadOnly(True)
        connect_layout.addWidget(self.connect_btn)
        connect_layout.addWidget(QLabel("可用资金:"))
        connect_layout.addWidget(self.capital_input)
        main_layout.addLayout(connect_layout)  # 新增连接控件

        # 导航按钮组 
        nav_btn_layout = QHBoxLayout()
        nav_btn_layout.setContentsMargins(10, 5, 10, 5)
        nav_btn_layout.setSpacing(15)
        
        self.auto_sl_btn = QPushButton("自动止盈止损")
        self.manual_trade_btn = QPushButton("人工交易")
        # 设置按钮样式
        for btn in [self.auto_sl_btn, self.manual_trade_btn]:
            btn.setCheckable(True)
            btn.setStyleSheet("""
                QPushButton {
                    padding: 8px 15px;
                    min-width: 120px;
                    background-color: #f0f0f0;
                    border: 1px solid #cccccc;
                    border-radius: 4px;
                }
                QPushButton:checked {
                    background-color: #0078d4;
                    color: white;
                }
            """)
        
        nav_btn_layout.addWidget(self.auto_sl_btn)
        nav_btn_layout.addStretch()
        nav_btn_layout.addWidget(self.manual_trade_btn)
        main_layout.addLayout(nav_btn_layout)

        # 堆叠窗口初始化 (关键修复：添加页面切换动画)
        self.stacked_widget = QStackedWidget()
        self.stacked_widget.setContentsMargins(0, 0, 0, 0)
        
        # 初始化所有页面
        self.pages = [
            self.init_auto_sl_ui(),  # 自动止盈止损页面
            self.init_manual_trade_ui()  # 新增人工交易页面
        ]
        for page in self.pages:
            self.stacked_widget.addWidget(page)
            
        main_layout.addWidget(self.stacked_widget)
        self.setCentralWidget(main_widget)

    def init_manual_trade_ui(self):
        """初始化人工交易页面 (恢复原有功能)"""
        widget = QWidget()
        layout = QVBoxLayout(widget)
        layout.setContentsMargins(15, 15, 15, 15)
        
        # 交易控件组
        trade_group = QHBoxLayout()
        self.symbol_input = QLineEdit(placeholderText="输入合约代码")
        self.direction_combo = QComboBox()
        self.direction_combo.addItems(["买入", "卖出"])
        self.quantity_input = QLineEdit("1")
        self.quantity_input.setValidator(QIntValidator(1, 100))
        self.trade_btn = QPushButton("执行交易")
        
        trade_group.addWidget(QLabel("合约:"))
        trade_group.addWidget(self.symbol_input)
        trade_group.addWidget(QLabel("方向:"))
        trade_group.addWidget(self.direction_combo)
        trade_group.addWidget(QLabel("数量:"))
        trade_group.addWidget(self.quantity_input)
        trade_group.addWidget(self.trade_btn)
        layout.addLayout(trade_group)
        
        # 持仓表格
        self.position_table = QTableWidget()
        self.position_table.setColumnCount(5)
        self.position_table.setHorizontalHeaderLabels(["合约", "方向", "数量", "均价", "盈亏"])
        self.position_table.horizontalHeader().setSectionResizeMode(QHeaderView.Stretch)
        layout.addWidget(self.position_table)
        
        return widget

    def _init_help_ui(self):
        """初始化帮助页面（补全实现）"""
        help_widget = QWidget()
        layout = QVBoxLayout(help_widget)
        layout.addWidget(QLabel("帮助文档\n\n1. 连接天勤账户后使用自动交易\n2. 人工交易需手动输入合约代码"))
        return help_widget  # 补充返回值

    def handle_page_change(self, index):
        """处理页面切换（修复重复定义问题）"""
        self.stacked_widget.setCurrentIndex(index)
        # 同步按钮状态
        self.auto_sl_btn.setChecked(index == 0)
        self.manual_trade_btn.setChecked(index == 1)

    def setup_connections(self):
        """建立信号连接（补充关键连接）"""
        # 连接账户按钮
        self.connect_btn.clicked.connect(self.connect_tq)
        
        # 连接页面切换信号
        self.auto_sl_btn.clicked.connect(lambda: self.handle_page_change(0))
        self.manual_trade_btn.clicked.connect(lambda: self.handle_page_change(1))
        
        # 连接资金更新信号
        self.update_signal.connect(self.capital_input.setText)

        # +++ 强制显示所有页面元素 +++
        self.stacked_widget.currentChanged.connect(lambda: self.stacked_widget.currentWidget().show())
        
        return main_widget  # 确保返回主控件

    def setup_connections(self):
        """建立信号连接（补充页面切换信号）"""
        self.auto_sl_btn.clicked.connect(lambda: self.handle_page_change(0))
        self.manual_trade_btn.clicked.connect(lambda: self.handle_page_change(1))
        # 添加其他必要的信号连接

    # 删除重复定义的 handle_page_change 方法

    def __init__(self):
        super().__init__()
        self.init_data()
        self.init_ui()     # 现在包含完整界面初始化
        self.setup_connections()
        self.setWindowTitle("资金管理系统")
        self.resize(1200, 800)
        self.handle_page_change(0)  # 默认显示第一个页面
        
        # 新增策略状态变量
        self.strategy_running = False
        self.open_strategy_running = False

if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = CapitalManager()
    window.show()
    sys.exit(app.exec_())